import re,os,sys,numpy as np
from scipy import stats
import matplotlib.pyplot as pyplot
import warnings
warnings.simplefilter('ignore', DeprecationWarning)

def updateframes(nfi,intervals):
    nfb = float(intervals[nfi][0]) if len(intervals) > nfi else float("infinity") #frame that starts next interval when p1,p2 are paired
    nfe = float(intervals[nfi][1]) if len(intervals) > nfi else float("infinity") #frame that starts next interval when p1,p2 are paired
    return nfb-1,nfe-1

def calcdts(pfile,basetracksdir):
    pdts = [] #array of distances from pairs
    npdts = [] #array of distances from non-pairs
    for line in pfile.readlines():
        line = line.strip()
        npair = []
        halves = line.split(':')
        framelist = halves[0].split(';')
        actors = halves[1].split(',')
        actors = (actors[0],actors[1]) if actors[0] < actors[1] else (actors[1],actors[0])
        pairings[actors] = []
        for frames in framelist:
            fpair = frames.split('-')
            f1 = fpair[0]
            f2 = fpair[1]
            pairings[actors].append((f1,f2))

    for p1file in sorted(os.listdir(basetracksdir)):
        if p1file.startswith("person_"):
            p1handle = open(basetracksdir+"/"+p1file,"r")
            p1frames = p1handle.readlines()
            p1coords = []
            for frame in p1frames:
                frame = frame.strip()
                coordlist = re.split("\s+",frame)
                p1coords.append((float(coordlist[0]),float(coordlist[2])) if coordlist[3].startswith("1") else (float("nan"),float("nan")))
            p1id = (((p1file.split("_"))[1]).split("."))[0] #grab id, which is between "_" and "." in file name
            for p2file in sorted(os.listdir(basetracksdir)):
                if p2file.startswith("person_"):
                    p2id = (((p2file.split("_"))[1]).split("."))[0]
                    if p2id <= p1id:
                        continue
                    p2handle = open(basetracksdir+"/"+p2file,"r")
                    p2frames = p2handle.readlines()
                    p2coords = []
                    for frame in p2frames:
                        frame = frame.strip()
                        coordlist = re.split("\s+",frame)
                        p2coords.append((float(coordlist[0]),float(coordlist[2])) if coordlist[3].startswith("1") else (float("nan"),float("nan")))
                    intervals = pairings[(p1id,p2id)] if (p1id,p2id) in pairings.keys() else []
                    nfi = 0
                    nfb,nfe = updateframes(nfi,intervals)
                    #print nfb,nfe
                    prevdist = float("nan")
                    for i,p1point in enumerate(p1coords):
                        p2point = p2coords[i]
                        if not np.isnan(p1point[0]) and not np.isnan(p2point[0]): #only want to include point if it isn't nan
                            xdist = np.fabs(p1point[0]-p2point[0])
                            zdist = np.fabs(p1point[1]-p2point[1])
                            curdist = np.sqrt(xdist*xdist+zdist*zdist)
                            if not np.isnan(prevdist):
                                dt = curdist-prevdist
                                #print i,nfb
                                if i >= nfb:
                                    #print "paired frame\n"
                                    if i > nfe:
                                        nfi += 1
                                        nfb,nfe = updateframes(nfi,intervals)
                                        #print nfb,nfe
                                        npdts.append(dt)
                                    pdts.append(dt)
                                else:
                                    npdts.append(dt)
                            prevdist = curdist
                        else:
                            prevdist = float("nan")
    return pdts,npdts


def calcdists(pfile,basetracksdir):
    pdists = [] #array of distances from pairs
    npdists = [] #array of distances from non-pairs
    for line in pfile.readlines():
        line = line.strip()
        npair = []
        halves = line.split(':')
        framelist = halves[0].split(';')
        actors = halves[1].split(',')
        actors = (actors[0],actors[1]) if actors[0] < actors[1] else (actors[1],actors[0])
        pairings[actors] = []
        for frames in framelist:
            fpair = frames.split('-')
            f1 = fpair[0]
            f2 = fpair[1]
            pairings[actors].append((f1,f2))

    for p1file in sorted(os.listdir(basetracksdir)):
        if p1file.startswith("person_"):
            p1handle = open(basetracksdir+"/"+p1file,"r")
            p1frames = p1handle.readlines()
            p1coords = []
            for frame in p1frames:
                frame = frame.strip()
                coordlist = re.split("\s+",frame)
                p1coords.append((float(coordlist[0]),float(coordlist[2])) if coordlist[3].startswith("1") else (float("nan"),float("nan")))
            p1id = (((p1file.split("_"))[1]).split("."))[0] #grab id, which is between "_" and "." in file name
            for p2file in sorted(os.listdir(basetracksdir)):
                if p2file.startswith("person_"):
                    p2id = (((p2file.split("_"))[1]).split("."))[0]
                    if p2id <= p1id:
                        continue
                    p2handle = open(basetracksdir+"/"+p2file,"r")
                    p2frames = p2handle.readlines()
                    p2coords = []
                    for frame in p2frames:
                        frame = frame.strip()
                        coordlist = re.split("\s+",frame)
                        p2coords.append((float(coordlist[0]),float(coordlist[2])) if coordlist[3].startswith("1") else (float("nan"),float("nan")))
                    intervals = pairings[(p1id,p2id)] if (p1id,p2id) in pairings.keys() else []
                    nfi = 0
                    nfb,nfe = updateframes(nfi,intervals)
                    #print nfb,nfe
                    for i,p1point in enumerate(p1coords):
                        p2point = p2coords[i]
                        if not np.isnan(p1point[0]) and not np.isnan(p2point[0]): #only want to include point if it isn't nan
                            xdist = np.fabs(p1point[0]-p2point[0])
                            zdist = np.fabs(p1point[1]-p2point[1])
                            curdist = np.sqrt(xdist*xdist+zdist*zdist)
                            #print i,nfb
                            if i >= nfb:
                                #print "paired frame\n"
                                if i > nfe:
                                    nfi += 1
                                    nfb,nfe = updateframes(nfi,intervals)
                                    #print nfb,nfe
                                    npdists.append(curdist)
                                pdists.append(curdist)
                            else:
                                npdists.append(curdist)
    return pdists,npdists

with np.errstate(invalid='ignore'):
    #runs from msave dir
    basetracksdir = "./data/"
    posdicts = dict()
    allddists = []
    pairings = dict()
    pdists = []
    npdists = []
    pdts = []
    npdts = []
    pdists = np.array(pdists)
    npdists = np.array(npdists)
    pdts = np.array(pdts)
    npdts = np.array(npdts)

    if len(sys.argv) > 1:
        basetracksdir += sys.argv[1]
        pfile = open(basetracksdir+"/.pairings","r")
        pdists,npdists = calcdists(pfile,basetracksdir)
        pdts,npdts = calcdts(pfile,basetracksdir)

    else: # if no sub-directory name passed as argument, run on all of them
        for subdir in sorted(os.listdir(basetracksdir)):
            curtracksdir = basetracksdir+subdir
            pfile = open(curtracksdir+"/.pairings","r")
            curpdists,curnpdists = calcdists(pfile,curtracksdir)
            curpdts,curnpdts = calcdts(pfile,curtracksdir)
            pdists = np.append(pdists,curpdists)
            npdists = np.append(npdists,curnpdists)
            pdts = np.append(pdts,curpdts)
            npdts = np.append(npdts,curnpdts)

    #palpha,pgloc,pbeta = stats.gamma.fit(pdists)
    #npalpha,npgloc,npbeta = stats.gamma.fit(npdists)
    ploc,pscale = stats.norm.fit(pdts,loc=0)
    nploc,npscale = stats.norm.fit(npdts,loc=0)
    pgloc,palpha = stats.expon.fit(pdists)
    npgloc,npalpha = stats.expon.fit(npdists)


    #print "p loc,alpha,beta,mean,sdev: ",pgloc,palpha,pbeta,(palpha*pbeta),(palpha*pbeta*pbeta)
    #print "np loc,alpha,beta,mean,sdev: ",npgloc,npalpha,npbeta,(npalpha*npbeta),(npalpha*npbeta*npbeta)
    print "p alpha,mean,sdev,lambda: ",palpha,(palpha),(palpha),1./palpha
    print "np alpha,mean,sdev,lambda: ",npalpha,(npalpha),(npalpha),1./npalpha
    print "pdt loc,scale: ",ploc,pscale
    print "npdt loc,scale: ",nploc,npscale

    pdistm = np.mean(pdists)
    npdistm = np.mean(npdists)
    pdtm = np.mean(pdts)
    npdtm = np.mean(npdts)
    pdtsd = np.std(pdts)
    npdtsd = np.std(npdts)

    print "p,np means: ",pdistm,npdistm
    print "pdt,npdt means,sdevs: ",pdtm,pdtsd,npdtm,npdtsd
    print stats.norm.stats(loc=ploc,scale=pscale,moments='mvsk')

    outfile = open('distparams','a')
    #outfile.write(str(palpha*pbeta)+"\n"+str(palpha*pbeta*pbeta)+"\n"+str(npalpha*npbeta)+"\n"+str(npalpha*npbeta*npbeta)+"\n"+str(npdtsd/pdtsd)+"\n")
    outfile.write(str(palpha)+"\n"+str(palpha)+"\n"+str(npalpha)+"\n"+str(npalpha)+"\n"+str(npdtsd/pdtsd)+"\n")
    outfile.close()
